// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: daniel

#include "coin2/exchange/bybit_common/api_util/rate_tracker.h"

#include <utility>

#include "coin2/exchange/base/market/market_exchange_api.h"

constexpr auto GET_METHOD_PER_IP = "get_method_per_ip";
constexpr auto POST_METHOD_PER_IP = "post_method_per_ip";
constexpr auto QUERY_PER_ENDPOINT = "query_per_endpoint";
constexpr auto PLACE_PER_ENDPOINT = "place_per_endpoint";
constexpr auto CANCEL_PER_ENDPOINT = "cancel_per_endpoint";
constexpr int64_t NS = 1'000'000'000LL;
constexpr auto EXAMPLE_SPOT_BYBIT = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke06/bybit/trade_key.json",
  "rate_limit_info": {
    "spot_bybit": {
      "rate_limits": {
        "get_method_per_ip": {
          "secs": 1,
          "reqs": 70
        },
        "post_method_per_ip": {
          "secs": 1,
          "reqs": 50
        },
        "query_per_endpoint": {
          "secs": 1,
          "reqs": 10
        },
        "place_per_endpoint": {
          "secs": 1,
          "reqs": 200
        },
        "cancel_per_endpoint": {
          "secs": 1,
          "reqs": 200
        }
      }
    }
  }
})";

constexpr auto EXAMPLE_FUTURES_BYBIT = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke06/bybit_futures/trade_key.json",
  "rate_limit_info": {
    "futures_bybit": {
      "rate_limits": {
        "get_method_per_ip": {
          "secs": 1,
          "reqs": 70
        },
        "post_method_per_ip": {
          "secs": 1,
          "reqs": 50
        },
        "query_per_endpoint": {
          "secs": 1,
          "reqs": 10
        },
        "place_per_endpoint": {
          "secs": 1,
          "reqs": 200
        },
        "cancel_per_endpoint": {
          "secs": 1,
          "reqs": 200
        }
      }
    }
  }
})";
namespace coin2::exchange::bybit_common::api_util {

namespace detail {

//////////////////////////////////////////////////////////////////////////////////
// class IPRateTracker
//////////////////////////////////////////////////////////////////////////////////
IPRateTracker::IPRateTracker(const coin2::exchange::order::RateLimitInfoFuturesBybit& cfg) {
  for (const auto& [key, val] : cfg.rate_limits()) {
    if (key == GET_METHOD_PER_IP) {
      get_method_rate_tracker_ = std::make_unique<FixedWindowRateTracker>(val.secs() * NS);
      get_method_rate_tracker_->SetRateLimitInfo(val.reqs(), val.reqs());
    } else if (key == POST_METHOD_PER_IP) {
      post_method_rate_tracker_ = std::make_unique<FixedWindowRateTracker>(val.secs() * NS);
      post_method_rate_tracker_->SetRateLimitInfo(val.reqs(), val.reqs());
    }
  }
  CHECK(get_method_rate_tracker_) << "missing get method rate tracker for Bybit";
  CHECK(post_method_rate_tracker_) << "missing post method rate tracker for Bybit";
}

bool IPRateTracker::HasEnoughRateLimit(int64_t ts, ActionType action_type) {
  switch (action_type) {
    case ActionType::QUERY:
      return get_method_rate_tracker_->GetRemaining(ts) > 10;  // add buffer to avoid ip banned
    case ActionType::PLACE:
    case ActionType::AMEND:
    case ActionType::CANCEL:
      return post_method_rate_tracker_->GetRemaining(ts) > 10;  // add buffer to avoid ip banned
    default:
      NOTREACHED() << "unsupported action type: " << action_type;
  }
}

void IPRateTracker::Push(int32_t num_reqs, int64_t timestamp, ActionType action_type) {
  switch (action_type) {
    case ActionType::QUERY:
      get_method_rate_tracker_->Push(num_reqs, timestamp);
      break;
    case ActionType::PLACE:
    case ActionType::AMEND:
    case ActionType::CANCEL:
      post_method_rate_tracker_->Push(num_reqs, timestamp);
      break;
    default:
      NOTREACHED() << "unsupported action type: " << action_type;
  }
}

void IPRateTracker::SetRateLimitInfo(int32_t remain, HttpMethod http_method) {
  switch (http_method) {
    case HttpMethod::GET: {
      auto assign = get_method_rate_tracker_->MaxReq();
      get_method_rate_tracker_->SetRateLimitInfo(assign, remain);
      break;
    }
    case HttpMethod::POST: {
      auto assign = post_method_rate_tracker_->MaxReq();
      post_method_rate_tracker_->SetRateLimitInfo(assign, remain);
      break;
    }
    default:
      NOTREACHED() << "unsupported method: " << http_method;
      break;
  }
}

auto IPRateTracker::GetMaxRateUsedRatio(int64_t ts) -> double {
  double get_max = 0.;
  if (get_method_rate_tracker_) {
    auto max_reqs = get_method_rate_tracker_->MaxReq();
    auto remaining = get_method_rate_tracker_->GetRemaining(ts);
    if (max_reqs < remaining) {
      // shouldn't reached here
      LOG(ERROR) << "ops~ something is wrong, rate_limit=" << max_reqs
                 << ", remaining=" << remaining;
      return 0.;
    }
    get_max = 1.0 * (max_reqs - remaining) / max_reqs;
  }
  double post_max = 0.;
  if (post_method_rate_tracker_) {
    auto max_reqs = post_method_rate_tracker_->MaxReq();
    auto remaining = post_method_rate_tracker_->GetRemaining(ts);
    if (max_reqs < remaining) {
      // shouldn't reached here
      LOG(ERROR) << "ops~ something is wrong, rate_limit=" << max_reqs
                 << ", remaining=" << remaining;
      return 0.;
    }
    post_max = 1.0 * (max_reqs - remaining) / max_reqs;
  }

  return std::max(get_max, post_max);
}

//////////////////////////////////////////////////////////////////////////////////
// class ApiRateTracker
//////////////////////////////////////////////////////////////////////////////////
ApiRateTracker::ApiRateTracker(
    const coin2::exchange::order::RateLimitInfoFuturesBybit& cfg,
    const std::string& rate_limit_category)
    : rate_limit_category_{rate_limit_category} {
  all_rate_limit_ = cfg;
  for (const auto& [key, val] : cfg.rate_limits()) {
    if (key == rate_limit_category_) {
      default_api_rate_limit_.MergeFrom(val);
      break;
    }
  }
  CHECK(default_api_rate_limit_.has_secs() && default_api_rate_limit_.has_reqs())
      << "missing api rate tracker for Bybit";
}

bool ApiRateTracker::HasEnoughRateLimit(int64_t ts, const std::string& endpoint) {
  if (rate_tracker_map_.count(endpoint) == 0) {
    auto api_rate_limit = default_api_rate_limit_;
    for (const auto& [key, val] : all_rate_limit_.rate_limits()) {
      if (key == endpoint) {
        api_rate_limit.MergeFrom(val);
        break;
      }
    }
    CHECK(api_rate_limit.secs() > 0) << "invalid api rate limit for endpoint: " << endpoint;
    CHECK(api_rate_limit.reqs() > 0) << "invalid api rate limit for endpoint: " << endpoint;
    rate_tracker_map_[endpoint] =
        std::make_unique<FixedWindowRateTracker>(api_rate_limit.secs() * NS);
    rate_tracker_map_[endpoint]->SetRateLimitInfo(api_rate_limit.reqs(), api_rate_limit.reqs());
    return true;
  }
  return rate_tracker_map_[endpoint]->GetRemaining(ts) > 0;
}

void ApiRateTracker::Push(int32_t num_reqs, int64_t ts, const std::string& endpoint) {
  if (rate_tracker_map_.count(endpoint) == 0) {
    return;
  }
  rate_tracker_map_[endpoint]->Push(num_reqs, ts);
}

void ApiRateTracker::SetRateLimitInfo(
    const std::string& endpoint,
    int32_t max_reqs,
    int32_t remain) {
  if (rate_tracker_map_.count(endpoint) == 0) {
    return;
  }
  auto& rate_limit_tracker = rate_tracker_map_[endpoint];
  auto assign = rate_limit_tracker->MaxReq();
  rate_limit_tracker->SetRateLimitInfo(assign, remain);
  if (assign != max_reqs) {
    LOG_EVERY_N(WARNING, 10) << "[" << endpoint << "] (" << assign << "|" << max_reqs << ")"
                             << "the max api rate limit is changed, please contact OG team";
  }
}

auto ApiRateTracker::GetMaxRateUsedRatio(int64_t ts, const std::string& endpoint) -> double {
  if (rate_tracker_map_.count(endpoint) == 0) {
    return 0.;
  }
  const auto& rate_tracker = rate_tracker_map_.at(endpoint);
  auto max_reqs = rate_tracker->MaxReq();
  auto remaining = rate_tracker->GetRemaining(ts);
  if (max_reqs < remaining) {
    // shouldn't reached here
    LOG(ERROR) << "ops~ something is wrong, rate_limit=" << max_reqs << ", remaining=" << remaining;
    return 0.;
  }
  return 1.0 * (max_reqs - remaining) / max_reqs;
}

}  // namespace detail

//////////////////////////////////////////////////////////////////////////////////
// class BybitRateTracker
//////////////////////////////////////////////////////////////////////////////////
BybitRateTracker::BybitRateTracker(
    const MarketExchangeApi& mea,
    const coin2::exchange::order::OrderConnectionConfig& cfg) {
  using RateConfig = coin2::exchange::base::order::RateConfig;
  RateConfig::Get()->GetConfig(mea, cfg, &rate_limit_info_);
  // validate config
  Validate(mea, rate_limit_info_);
  // init rate tracker
  query_api_rate_tracker_ =
      std::make_unique<detail::ApiRateTracker>(rate_limit_info_, QUERY_PER_ENDPOINT);
  place_api_rate_tracker_ =
      std::make_unique<detail::ApiRateTracker>(rate_limit_info_, PLACE_PER_ENDPOINT);
  cancel_api_rate_tracker_ =
      std::make_unique<detail::ApiRateTracker>(rate_limit_info_, CANCEL_PER_ENDPOINT);
  // rate limit report
  rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
  rate_limit_report_writer_->SetRateLimitInfoString(&rate_limit_info_);
}

void BybitRateTracker::Validate(
    const MarketExchangeApi& mea,
    const coin2::exchange::order::RateLimitInfoFuturesBybit& rate_limit_info) const {
  std::string_view example_view =
      mea.market == coin::proto::MarketType::Spot ? EXAMPLE_SPOT_BYBIT : EXAMPLE_FUTURES_BYBIT;
  // check mandatory field
  if (rate_limit_info.rate_limits().count(GET_METHOD_PER_IP) == 0 ||
      rate_limit_info.rate_limits().count(POST_METHOD_PER_IP) == 0 ||
      rate_limit_info.rate_limits().count(QUERY_PER_ENDPOINT) == 0 ||
      rate_limit_info.rate_limits().count(PLACE_PER_ENDPOINT) == 0 ||
      rate_limit_info.rate_limits().count(CANCEL_PER_ENDPOINT) == 0) {
    using RateConfig = coin2::exchange::base::order::RateConfig;
    RateConfig::CheckFail(mea.String(), example_view);
    CHECK(false) << "invalid rate limit config for Bybit";
  }
  // print rate limit info
  std::stringstream ss;
  ss << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
     << "\nOG request rate was set as below:\n";
  for (auto [key, val] : rate_limit_info.rate_limits()) {
    ss << "\n" << key << ": [" << val.reqs() << " (reqs) / " << val.secs() << " (seconds)]";
  }
  ss << "\n\nOr you can configure customized rate in driver config. For example: " << example_view
     << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";
  LOG(INFO) << ss.str();
}

double BybitRateTracker::GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) {
  const std::string bind_ip = "";  // TODO(daniel): handle bind ip case
  double ret = 0.;
  switch (action_type) {
    case ActionType::QUERY:
      break;
    case ActionType::PLACE: {
      double api_max =
          place_api_rate_tracker_->GetMaxRateUsedRatio(timestamp, submit_order_endpoint_);
      double ip_max = 0.;
      if (ip_rate_tracker_map_.count(bind_ip)) {
        ip_max = ip_rate_tracker_map_[bind_ip]->GetMaxRateUsedRatio(timestamp);
      }
      ret = std::max(api_max, ip_max);
      break;
    }
    case ActionType::AMEND:
      break;
    case ActionType::CANCEL: {
      double api_max =
          cancel_api_rate_tracker_->GetMaxRateUsedRatio(timestamp, cancel_order_endpoint_);
      double ip_max = 0.;
      if (ip_rate_tracker_map_.count(bind_ip)) {
        ip_max = ip_rate_tracker_map_[bind_ip]->GetMaxRateUsedRatio(timestamp);
      }
      ret = std::max(api_max, ip_max);
      break;
    }
    default:
      NOTREACHED() << "unsupported action type: " << action_type;
      break;
  }
  return ret;
}

void BybitRateTracker::SetRateLimitInfoString(const google::protobuf::Message* message) {
  rate_limit_report_writer_->SetRateLimitInfoString(message);
}

void BybitRateTracker::WriteRateLimitReportMessage(OrderLoggerCB callback) {
  if (!callback) {
    return;
  }
  rate_limit_report_writer_->Write(callback);
}

void BybitRateTracker::UpdateReqRateInfo(const nlohmann::json& info) {
  if (0 == info["X-Bip-Get-Limit-Status"].get<int>()) {
    const auto& bind_ip = info["bind_ip"].get<std::string>();
    if (ip_rate_tracker_map_.count(bind_ip)) {
      ip_rate_tracker_map_[bind_ip]->SetRateLimitInfo(0, HttpMethod::GET);
    }
  }
  if (0 == info["X-Bip-Post-Limit-Status"].get<int>()) {
    const auto& bind_ip = info["bind_ip"].get<std::string>();
    if (ip_rate_tracker_map_.count(bind_ip)) {
      ip_rate_tracker_map_[bind_ip]->SetRateLimitInfo(0, HttpMethod::POST);
    }
  }
  const auto& remain = info["X-Bapi-Limit-Status"].get<int>();
  if (0 <= remain) {
    const auto& endpoint = info["X-Bapi-Limit-Endpoint"].get<std::string>();
    const auto& max_reqs = info["X-Bapi-Limit"].get<int>();
    if (endpoint == submit_order_endpoint_) {
      place_api_rate_tracker_->SetRateLimitInfo(endpoint, max_reqs, remain);
    } else if (endpoint == cancel_order_endpoint_) {
      cancel_api_rate_tracker_->SetRateLimitInfo(endpoint, max_reqs, remain);
    } else {
      query_api_rate_tracker_->SetRateLimitInfo(endpoint, max_reqs, remain);
    }
  }
  if (0 == info["X-Bapi-Limit-Status"].get<int>()) {
    int64_t five_secs_later = GetCurrentTimestamp() + 5'000'000'000LL;
    int64_t retry_after = info["X-Bapi-Limit-Reset-Timestamp"].get<int64_t>();
    retry_after *= 1000'000LL;
    if (retry_after > 0 && retry_after < five_secs_later) {
      const auto& endpoint = info["X-Bapi-Limit-Endpoint"].get<std::string>();
      const auto& bind_ip = info["bind_ip"].get<std::string>();
      this->BanApiForAPeriod(BanType::SPECIFIC_API, endpoint, bind_ip, retry_after);
    }
  }
  if (0 == info["X-Bip-Get-Limit-Status"].get<int>()) {
    const auto& endpoint = info["X-Bapi-Limit-Endpoint"].get<std::string>();
    const auto& bind_ip = info["bind_ip"].get<std::string>();
    const int64_t retry_after = GetCurrentTimestamp() + 30 * 60 * 1'000'000'000LL /* 30 mins */;
    this->BanApiForAPeriod(BanType::GET_API, endpoint, bind_ip, retry_after);
  }
  if (0 == info["X-Bip-Post-Limit-Status"].get<int>()) {
    const auto& endpoint = info["X-Bapi-Limit-Endpoint"].get<std::string>();
    const auto& bind_ip = info["bind_ip"].get<std::string>();
    const int64_t retry_after = GetCurrentTimestamp() + 30 * 60 * 1'000'000'000LL /* 30 mins */;
    this->BanApiForAPeriod(BanType::POST_API, endpoint, bind_ip, retry_after);
  }
}

bool BybitRateTracker::HasEnoughRateLimit(
    int64_t ts,
    ActionType action_type,
    const std::string& endpoint,
    const std::string& bind_ip) {
  if (!IsBannedPeriod(ts, endpoint, bind_ip) &&
      CheckRateLimit(ts, action_type, endpoint, bind_ip)) {
    return true;
  }

  rate_limit_report_writer_->SetHitRateLimit(ts, action_type);
  return false;
}

void BybitRateTracker::Push(
    int32_t num_reqs,
    int64_t ts,
    ActionType action_type,
    const std::string& endpoint,
    const std::string& bind_ip) {
  // update ip rate limit
  if (ip_rate_tracker_map_.count(bind_ip) == 0) {
    ip_rate_tracker_map_[bind_ip] = std::make_unique<detail::IPRateTracker>(rate_limit_info_);
  }
  ip_rate_tracker_map_[bind_ip]->Push(num_reqs, ts, action_type);
  // update symbol rate limit
  switch (action_type) {
    case ActionType::QUERY:
      query_api_rate_tracker_->Push(num_reqs, ts, endpoint);
      break;
    case ActionType::PLACE:
      place_api_rate_tracker_->Push(num_reqs, ts, endpoint);
      break;
    case ActionType::AMEND:
      // ignored
      break;
    case ActionType::CANCEL:
      cancel_api_rate_tracker_->Push(num_reqs, ts, endpoint);
      break;
    default:
      break;
  }
}

bool BybitRateTracker::CheckRateLimit(
    int64_t ts,
    ActionType action_type,
    const std::string& endpoint,
    const std::string& bind_ip) {
  bool ret = true;
  // check api rate limit
  switch (action_type) {
    case ActionType::QUERY:
      ret = query_api_rate_tracker_->HasEnoughRateLimit(ts, endpoint);
      break;
    case ActionType::PLACE:
      ret = place_api_rate_tracker_->HasEnoughRateLimit(ts, endpoint);
      break;
    case ActionType::AMEND:
      // ignored
      break;
    case ActionType::CANCEL:
      ret = cancel_api_rate_tracker_->HasEnoughRateLimit(ts, endpoint);
      break;
    default:
      break;
  }
  if (ret) {
    // check ip rate limit
    if (ip_rate_tracker_map_.count(bind_ip) == 0) {
      ip_rate_tracker_map_[bind_ip] = std::make_unique<detail::IPRateTracker>(rate_limit_info_);
    }
    ret = ip_rate_tracker_map_[bind_ip]->HasEnoughRateLimit(ts, action_type);
  }

  return ret;
}

bool BybitRateTracker::IsBannedPeriod(
    int64_t timestamp,
    const std::string& endpoint,
    const std::string& bind_ip) {
  const std::string key = fmt::format("{}{}", bind_ip, endpoint);
  if (enable_time_by_ip_map_.count(key) == 0) {
    return false;
  }

  const auto& enable_time = enable_time_by_ip_map_.at(key);

  if (timestamp >= enable_time) {
    return false;
  }

  LOG_EVERY_N(WARNING, 50) << "Hit api temporary disabled until " << enable_time << ", " << endpoint
                           << ", " << bind_ip;
  return true;
}

void BybitRateTracker::BanApiForAPeriod(
    BanType ban_type,
    const std::string& endpoint,
    const std::string& bind_ip,
    int64_t retry_after) {
  switch (ban_type) {
    case BanType::SPECIFIC_API: {
      const std::string key = fmt::format("{}{}", bind_ip, endpoint);
      enable_time_by_ip_map_[key] = retry_after;
      LOG(INFO) << "[Bybit] api rate limit hit, banning " << key << " till " << retry_after;
      break;
    }
    case BanType::GET_API: {
      for (const auto& ep : http_get_endpoints_) {
        const std::string key = fmt::format("{}{}", bind_ip, ep);
        enable_time_by_ip_map_[key] = retry_after;
        LOG(INFO) << "[Bybit] ip rate limit hit, banning " << key << " till " << retry_after;
      }
      break;
    }
    case BanType::POST_API: {
      for (const auto& ep : http_post_endpoints_) {
        const std::string key = fmt::format("{}{}", bind_ip, ep);
        enable_time_by_ip_map_[key] = retry_after;
        LOG(INFO) << "[Bybit] ip rate limit hit, banning " << key << " till " << retry_after;
      }
      break;
    }
    default:
      NOTREACHED() << "unsupported type: " << ban_type;
      break;
  }
}

}  // namespace coin2::exchange::bybit_common::api_util
